package com.menghao.rpc.netty;

import com.menghao.rpc.RpcConstants;
import com.menghao.rpc.netty.model.HostKey;
import com.menghao.rpc.netty.model.TcpConnection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * <p>TcpConnection容器.<br>
 * <p>使用读写锁保证线程安全</p>
 * <p>保证指向一台机器的连接（长连接）唯一</p>
 *
 * @author MarvelCode.
 */
public class TcpConnectionContainer {

    private static final Logger LOG = LoggerFactory.getLogger(TcpConnectionContainer.class);

    private Map<HostKey, TcpConnection> connectionMap = new HashMap<HostKey, TcpConnection>();

    private ReadWriteLock lock = new ReentrantReadWriteLock();

    public void register(String ip, int port) {
        HostKey hostKey = new HostKey(ip, port);
        if (connectionMap.get(hostKey) == null) {
            lock.writeLock().lock();
            try {
                if (connectionMap.get(hostKey) == null) {
                    try {
                        LOG.info(RpcConstants.LOG_RPC_PREFIX + "create ip-{},port-{} connection", ip, port);
                        connectionMap.put(hostKey, new TcpConnection(ip, port));
                    } catch (InterruptedException e) {
                        LOG.error(RpcConstants.LOG_RPC_PREFIX + "connect ip-{} port-{} exception,", e, ip, port);
                    }
                }
            } finally {
                lock.writeLock().unlock();
            }
        }
    }

    public TcpConnection get(String ip, int port) {
        lock.readLock().lock();
        try {
            HostKey hostKey = new HostKey(ip, port);
            return connectionMap.get(hostKey);
        } finally {
            lock.readLock().unlock();
        }
    }

    public void remove(String ip, int port) {
        HostKey hostKey = new HostKey(ip, port);
        if (connectionMap.get(hostKey) != null) {
            lock.writeLock().lock();
            try {
                if (connectionMap.get(hostKey) != null) {
                    TcpConnection connection = connectionMap.get(hostKey);
                    LOG.info(RpcConstants.LOG_RPC_PREFIX + "lost ip-{},port-{} connection", ip, port);
                    if (connection.isActive()) {
                        connection.close();
                    }
                    connectionMap.remove(hostKey);
                }
            } finally {
                lock.writeLock().unlock();
            }
        }
    }
}
